package com.wan.net.rpc.registry;

import com.wan.core.shutdown.ShutdownManager;
import com.wan.middleware.zookeeper.ZkConfig;
import com.wan.middleware.zookeeper.ZookeeperClient;
import com.wan.net.rpc.service.RpcService;
import com.wan.util.serializer.ProtoStuffSerializer;
import org.apache.curator.framework.recipes.cache.CuratorCacheListener;
import org.apache.zookeeper.CreateMode;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

/**
 * @author wwh
 * @date 2022/6/21 15:45
 */
public class ZookeeperRegistry extends Registry.LocalRegistry {
    private static final Logger logger = LoggerFactory.getLogger(ZookeeperRegistry.class);
    private static final String ROOT = "/rpcServices";

    private final ZookeeperClient client;

    public ZookeeperRegistry(String address) throws Exception {
        ZkConfig zkConfig = new ZkConfig();
        zkConfig.setAddress(address);
        this.client = new ZookeeperClient(zkConfig);
        initCache();
        ShutdownManager.addShutdown(client::close);
    }

    private void initCache() throws Exception {
        //根目录创建
        client.createPathIfNotExist(ROOT);
        CountDownLatch countDownLatch = new CountDownLatch(1);
        //监听变化
        CuratorCacheListener curatorCacheListener = CuratorCacheListener.builder()
            .forCreates(data -> {
                try {
                    if (data.getPath().equals(ROOT)) {
                        return;
                    }
                    RpcService service = ProtoStuffSerializer.decode(data.getData(), RpcService.class);
                    super.registry(service);
                } catch (Exception e) {
                    logger.error("zookeeper创建节点异常,node:{}", data.getPath(), e);
                }
            })
            .forDeletes(data -> {
                try {
                    RpcService delService = ProtoStuffSerializer.decode(data.getData(), RpcService.class);
                    super.unregister(delService);
                } catch (Exception e) {
                    logger.error("zookeeper删除节点异常,node:{}", data.getPath(), e);
                }
            })
            .forInitialized(countDownLatch::countDown).build();
        client.watchCuratorCache(ROOT,curatorCacheListener);
        if(!countDownLatch.await(3,TimeUnit.SECONDS)){
            logger.error("zookeeper curatorCache initialized failed");
        }
    }

    @Override
    public void registry(RpcService rpcService) {
        try {
            String serviceName = rpcService.uniqueMD5();
            String path = ROOT + "/" + serviceName;
            byte[] data = ProtoStuffSerializer.encode(rpcService);
            String pathData = client.createPathData(path, data, CreateMode.PERSISTENT);
            logger.info("[rpc]服务注册成功,path:{},service:{}",pathData,rpcService);
        }catch (Exception e){
            logger.error("[rpc]服务注册失败,service:{}",rpcService,e);
            throw new RuntimeException(e);
        }
    }

    @Override
    public void unregister(RpcService rpcService) {
        try {
            String serviceName = rpcService.uniqueMD5();
            String path = ROOT + "/" + serviceName;
            client.deletePath(path);
        }catch (Exception e){
            logger.error("[rpc]服务删除失败,service:{}",rpcService,e);
            throw new RuntimeException(e);
        }
    }
}
